from volcenginesdkarkruntime import Ark
from flask import Flask, request
from flask_cors import CORS
from io import BytesIO
from PIL import Image
import base64
import re


app = Flask(__name__)
cors = CORS()
cors.init_app(app)

client = Ark(
    api_key='请填写自己的API密钥',
    base_url="https://ark.cn-beijing.volces.com/api/v3",
)


@app.route('/ai', methods=['POST'])
def ai_check():
    # 获取到的数据格式：[ '0-4b692f', '63-4b692f' ]
    data = request.get_json()
    if not data:
        return {
            'status': 'fail',
            'msg': 'Fail to receive data.'
        }

    try:
        img = draw(data)
        base64_str = encode_image(img)
        isOk, check_result = ask_doubao('请鉴定下这幅8x8大小的武器像素画，并打个分数，满分200分。', base64_str)
        if not isOk:
            return {
                'status': 'fail',
                'msg': 'Fail to receive data from ai.'
            }

        return {
            'status': 'success',
            'msg': check_result
        }
    except Exception as e:
        return {
            'status': 'fail',
            'msg': str(e)
        }


def draw(data=None):
    # 初始化画布大小和每个像素点大小
    img = Image.new('RGB', (160, 160), 15066597)
    item_size = 20

    pixels = img.load()
    
    for pixel_info in data:
        # 像素点位置
        pixel_index = int(pixel_info.split('-')[0])
        y = int(pixel_index / 8)
        x = pixel_index - y*8

        # 像素点颜色
        pixel_color = []
        for value in pixel_info.split('-')[1].split(','):
            pixel_color.append(int(value))
        pixel_color = tuple(pixel_color)

        # 绘制
        for i in range(x*item_size, x*item_size+20):
            for j in range(y*item_size, y*item_size+20):
                pixels[i, j] = pixel_color

    return img


def encode_image(img):
    buffered = BytesIO()  
    img.save(buffered, format='PNG')  
    img_byte_arr = buffered.getvalue()  
    base64_str = base64.b64encode(img_byte_arr).decode('utf-8')  
    return base64_str


def ask_doubao(text, imgBase64):
    response = client.chat.completions.create(
        model="ep-20250115164341-d97rh",
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": f"{text}"},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/png;base64,{imgBase64}"
                        }
                    },
                ],
            }
        ],

        # 免费开启推理会话应用层加密，访问 https://www.volcengine.com/docs/82379/1389905 了解更多
        extra_headers={'x-is-encrypted': 'true'},
    )

    if not response:
        return False, None
    
    content = response.choices[0].message.content
    score_list = re.findall('[0-9]+\s?分', content)
    final_score = 0
    if score_list:
        final_score = int(max(score_list, key=lambda x:int(x.replace('分', '').strip())).replace('分', '').strip())

    print(content)
    print(score_list)
    print(final_score)
   
    return True, {
        'score': final_score,
        'content': content.strip()
    }


if __name__ == '__main__':
    app.run(host='0.0.0.0')